#!/usr/bin/env python3
from bcc import BPF
from collections import defaultdict

# BPF 程序
bpf_program = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
#include <linux/fs.h>

struct data_t {
    u32 pid;
    char comm[TASK_COMM_LEN];
    char filename[DNAME_INLINE_LEN];
};

BPF_HASH(counts, struct data_t);

int trace_write_return(struct pt_regs *ctx, struct file *file) {
    struct data_t data = {};
    u32 pid = bpf_get_current_pid_tgid();
    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    data.pid = pid;

    struct dentry *de = file->f_path.dentry;
    bpf_probe_read_kernel(&data.filename, sizeof(data.filename), de->d_name.name);

    counts.increment(data);
    return 0;
}
"""

# 加载 BPF 程序
b = BPF(text=bpf_program)

# 在内核函数 `vfs_write` 的返回位置附加 BPF 程序
b.attach_kretprobe(event="vfs_write", fn_name="trace_write_return")

print("Tracing file writes... Hit Ctrl-C to end.")

counts = defaultdict(int)

# 捕获输出
try:
    while True:
        # 读取 BPF 哈希表中的数据
        for k, v in b["counts"].items():
            counts[(k.comm.decode('utf-8', 'replace'), k.filename.decode('utf-8', 'replace'))] += v.value
        b["counts"].clear()
except KeyboardInterrupt:
    pass

# 打印统计结果
print("\n文件写入次数统计:")
for (comm, filename), count in sorted(counts.items(), key=lambda x: x[1], reverse=True):
    print(f"{comm} 写入 {filename}: {count} 次")